import os
from docx import Document

def read_and_number_headings(doc_path):
    # 确保文件存在
    if not os.path.exists(doc_path):
        print(f"ERROR: 文件未找到 -> {doc_path}")
        return {}

    # 读取文档
    doc = Document(doc_path)
    print(f"成功打开文件 -> {doc_path}")

    # 存储每一层级的编号
    heading_numbers = []

    # 存储标题和对应内容的字典
    headings_dict = {}

    print("开始遍历标题段落...")

    last_number = None  # 记录最后一次处理的编号，避免出现跳跃

    # 遍历段落，检查每个段落的样式是否为标题样式
    for para in doc.paragraphs:
        text = para.text.strip()
        style_name = para.style.name
        
        # 只处理包含"Heading"的段落（标题段落）
        if "Heading" in style_name:
            # 获取当前层级
            heading_level = int(style_name.split(' ')[-1])  # 获取 Heading 级别，例如 Heading 1 -> 1, Heading 2 -> 2
            
            # 如果当前层级大于现有的层级，说明是一个子标题，需扩展编号
            if heading_level > len(heading_numbers):
                for i in range(heading_level-len(heading_numbers)):
                    heading_numbers.append(1)  # 新增层级从1开始
            elif heading_level < len(heading_numbers):
                # 如果当前层级小于之前的层级，删除尾部多余的部分
                heading_numbers = heading_numbers[:heading_level]

                # 更新当前层级的编号
                heading_numbers[heading_level - 1] += 1

            else:
                # 如果当前层级等于现有层级，递增该层级的编号
                heading_numbers[heading_level - 1] += 1

            # 生成编号内容
            heading_number = '.'.join(map(str, heading_numbers))
            
            # 如果编号与上一行相同，表示跳跃编号，跳过当前段落
            if last_number == heading_number:
                continue

            last_number = heading_number  # 更新上一编号
            
            # 保存标题和对应内容到字典
            if heading_number not in headings_dict:
                headings_dict[heading_number] = {"title": text, "content": []}
            else:
                headings_dict[heading_number]["title"] = text

            # 打印输出，标题对齐
            print(f"{heading_number:<10} {text}")
        
        else:
            # 将非标题段落内添加到字典对应标题下
            if heading_numbers:
                heading_number = '.'.join(map(str, heading_numbers))
                headings_dict[heading_number]["content"].append(text)
    
    return headings_dict

def search_heading_by_number(headings_dict, heading_number):
    # 查找指定标题编号的内容
    if heading_number in headings_dict:
        heading = headings_dict[heading_number]
        print(f"查找的标题编号: {heading_number}")
        print(f"标题内容: {heading['title']}")
        print(f"段落内容: {heading['content']}")
    else:
        print(f"未找到编号为 {heading_number} 的标题内容。")

# 运行代码，替换成你的文件路径
doc_path = "/Users/guhaoyu/Desktop/造价AI工作/2、客户提供的建设需求或可研文档/01.docx"
headings_dict = read_and_number_headings(doc_path)

# 示例查找功能
search_heading_by_number(headings_dict, "1.1.1.6.3")  # 替换成你需要查找的标题编号
